#include "../Base.h"

NTSTATUS ReadMemoryDirect(PEPROCESS process,PVOID targetAddress, PVOID buffer, SIZE_T size) {
    
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    KAPC_STATE apc_state = { 0 };
    PVOID mem = ExAllocatePoolWithTag(NonPagedPool, size, 'yang');
    memset(mem, 0, size);
    KeStackAttachProcess(process, &apc_state);
    if (MmIsAddressValid(targetAddress) && MmIsAddressValid((PVOID)((ULONG64)targetAddress + size))) {
        memcpy(mem, targetAddress, size);
        status = STATUS_SUCCESS;
    }
    KeUnstackDetachProcess(&apc_state);
    if (NT_SUCCESS(status)) {
        memcpy(buffer, mem, size);
    }
    ExFreePoolWithTag(mem, 'yang');
    return status;
}


NTSTATUS ReadMemoryByIntelFunc(PEPROCESS process,PVOID targetAddress, PVOID buffer, SIZE_T size) {

    SIZE_T retSize;
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    status = MmCopyVirtualMemory(process, targetAddress, IoGetCurrentProcess(), buffer, size, UserMode, &retSize);
    return status;
}



PVOID MdlMapMemory(PMDL* mdl, PVOID targetAddress, SIZE_T size, MODE preMode) {
    PMDL pMdl = IoAllocateMdl(targetAddress, size, FALSE, FALSE,NULL);
    *mdl = pMdl;
    PVOID mapAddr = NULL;
    BOOLEAN isLock = FALSE;
    if (!pMdl) {
        return NULL;
    }
    __try {

        MmProbeAndLockPages(pMdl, KernelMode, IoReadAccess);
        isLock = TRUE;
        mapAddr = MmMapLockedPagesSpecifyCache(pMdl, KernelMode, MmNonCached, NULL, FALSE, NormalPagePriority);
    }
    __except (1) {

        if (isLock) {
            MmUnlockPages(pMdl);
        }
        IoFreeMdl(pMdl);
    }

    return mapAddr;

    
    
}

VOID MdlUnMapMemory(PMDL mdl, PVOID mapAddr) {

    BOOLEAN lock = FALSE;
    __try {

        MmUnmapLockedPages(mapAddr, mdl);
        lock = TRUE;
        MmUnlockPages(mdl);
        IoFreeMdl(mdl);
    }
    __except (1) {

        return;
    }
}


NTSTATUS ReadMemoryByMdl(PEPROCESS process, PVOID targetAddress, PVOID buffer, SIZE_T size) {

    
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    KAPC_STATE apc_state = { 0 };
    PVOID mem = ExAllocatePoolWithTag(NonPagedPool, size, 'yang');
    memset(mem, 0, size);
    KeStackAttachProcess(process, &apc_state);
    PMDL mdl= NULL;
    PVOID map = MdlMapMemory(&mdl, targetAddress, size, UserMode);
    if (map) {
        memcpy(mem, map, size);
        MdlUnMapMemory(mdl, map);
        status = STATUS_SUCCESS;
    }
    KeUnstackDetachProcess(&apc_state);
    if (NT_SUCCESS(status)) {
        memcpy(buffer, mem, size);
    }
    ExFreePoolWithTag(mem, 'yang');
 
    return status;

}

NTSTATUS ReadMemoryByCr3(PEPROCESS process, PVOID targetAddress, PVOID buffer, SIZE_T size) {


    NTSTATUS status = STATUS_UNSUCCESSFUL;
    KAPC_STATE apc_state = { 0 };
    PVOID mem = ExAllocatePoolWithTag(NonPagedPool, size, 'yang');
    memset(mem, 0, size);
    ULONG64 newCr3 = *(PULONG64)((PUCHAR)process + 0x28);
    ULONG64 oldCr3 = __readcr3();
    KeEnterCriticalRegion();
    _disable();
    __writecr3(newCr3);
    if (MmIsAddressValid(targetAddress) && MmIsAddressValid((PVOID)((ULONG64)targetAddress + size))) {
        memcpy(mem, targetAddress, size);
        status = STATUS_SUCCESS;
    }
    __writecr3(oldCr3);
    _enable();
    KeLeaveCriticalRegion();
    if (NT_SUCCESS(status)) {
        memcpy(buffer, mem, size);
    }
    ExFreePoolWithTag(mem, 'yang');
    return status;


}


NTSTATUS ReadMmeory(HANDLE pid, PVOID targetAddress, PVOID buffer, SIZE_T size, READWRITEMETHOD method)
{
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    if (targetAddress >= MmHighestUserAddress 
        || ((PVOID)((PUCHAR)targetAddress + size)) >= MmHighestUserAddress
        || ((ULONG64)((PUCHAR)targetAddress + size)) < (ULONG64)targetAddress) {
        return STATUS_ACCESS_VIOLATION;
    }

    if (buffer == NULL)return STATUS_INVALID_PARAMETER_3;
    PEPROCESS process;

    status = PsLookupProcessByProcessId(pid, &process);
    if (!NT_SUCCESS(status)) {
        return status;
    }
    if (PsGetProcessExitStatus(process) != 0x103) {
        ObDereferenceObject(process);
        return STATUS_INVALID_PARAMETER_1;
    }
    switch (method)
    {
    case READWRITEATTACH: {

        status = ReadMemoryDirect(process, targetAddress, buffer, size);
        break;

    }
    case READWRITECR3: {
        status = ReadMemoryByCr3(process, targetAddress, buffer, size);
        break;
    }
    case READWRITEINTEL: {
        status = ReadMemoryByIntelFunc(process, targetAddress, buffer, size);
        break;
    }
    case READWRITEMDL: {
        status = ReadMemoryByMdl(process, targetAddress, buffer, size);
        break;
    }
    default:
        break;
    }
   

    ObDereferenceObject(process);
    return status;
}

NTSTATUS WriteMemory(HANDLE pid, PVOID targetAddress, PVOID buffer, SIZE_T size)
{
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    if (targetAddress >= MmHighestUserAddress
        || ((PVOID)((PUCHAR)targetAddress + size)) >= MmHighestUserAddress
        || ((ULONG64)((PUCHAR)targetAddress + size)) < (ULONG64)targetAddress) {
        return STATUS_ACCESS_VIOLATION;
    }

    if (buffer == NULL)return STATUS_INVALID_PARAMETER_3;
    PEPROCESS process,srcProcess;
    srcProcess  = IoGetCurrentProcess();
    status = PsLookupProcessByProcessId(pid, &process);
    if (!NT_SUCCESS(status)) {
        return status;
    }
    if (PsGetProcessExitStatus(process) != 0x103) {
        ObDereferenceObject(process);
        return STATUS_INVALID_PARAMETER_1;
    }
    

    SIZE_T retSize = 0;
    status = MmCopyVirtualMemory(IoGetCurrentProcess(), buffer, process, targetAddress, size, UserMode, &retSize);
    if (NT_SUCCESS(status)) {
        ObDereferenceObject(process);
        return STATUS_SUCCESS;
    }

    KAPC_STATE apc_state;
    KeStackAttachProcess(process, &apc_state);

    PVOID baseAddr = targetAddress;
    SIZE_T tempSize = size;
    ULONG proAttr = 0;
    status = NtProtectVirtualMemory(NtCurrentProcess(), &baseAddr,&tempSize,PAGE_EXECUTE_READWRITE,&proAttr );
    if (NT_SUCCESS(status)) {
        retSize = 0;
        status = MmCopyVirtualMemory(srcProcess, buffer, process, targetAddress, size, UserMode, &retSize);
        NtProtectVirtualMemory(NtCurrentProcess(), &baseAddr, &tempSize, proAttr, &proAttr);
    }
    KeUnstackDetachProcess(&apc_state);

    if (NT_SUCCESS(status)) {
        ObDereferenceObject(process);
        return STATUS_SUCCESS;
    }

    ULONG64 oldCr0 = wpoff();
    retSize = 0;
    status = MmCopyVirtualMemory(srcProcess, buffer, process, targetAddress, size, UserMode, &retSize);
    NtProtectVirtualMemory(NtCurrentProcess(), &baseAddr, &tempSize, proAttr, &proAttr);
    wpOn(oldCr0);
   
    ObDereferenceObject(process);
    return status;
}

NTSTATUS QueryMemory(HANDLE pid, PVOID targetAddress, PMEMORY_BASIC_INFORMATION memoryBasicInfo)
{

    NTSTATUS status = STATUS_UNSUCCESSFUL;
    if (!memoryBasicInfo) return status;

    PEPROCESS process;
    status = PsLookupProcessByProcessId(pid, &process);
    if (!NT_SUCCESS(status)) {
        return status;
    }
    if (PsGetProcessExitStatus(process) != 0x103) {
        ObDereferenceObject(process);
        return STATUS_INVALID_PARAMETER_1;
    }
    
    PMEMORY_BASIC_INFORMATION mem = (PMEMORY_BASIC_INFORMATION)ExAllocatePoolWithTag(NonPagedPool, sizeof(MEMORY_BASIC_INFORMATION),'yang');
    RtlZeroMemory(mem, sizeof(MEMORY_BASIC_INFORMATION));
    KAPC_STATE apc_state;
    KeStackAttachProcess(process, &apc_state);
    SIZE_T retSize = 0;
    status = ZwQueryVirtualMemory(NtCurrentProcess(), targetAddress, MemoryBasicInformation, mem, sizeof(MEMORY_BASIC_INFORMATION), &retSize);

    KeUnstackDetachProcess(&apc_state);
    if (NT_SUCCESS(status)) {
        memcpy(memoryBasicInfo,mem , sizeof(MEMORY_BASIC_INFORMATION));
    }
    ExFreePoolWithTag(mem, 'yang');
    return status;
}
